In [1]:
# HMM - neprekidne raspodele emisija oblika: b = { 'SYMBOL': (MEAN, STD) }
# Logaritmovane vrednosti raspodela
class HMM:
    def __init__(self, l = None):        
        if l != None:
            A, B, P = l

            # log(A)
            logA = {}
            for qi in A:
                if qi not in logA:
                    logA[qi] = {}
                for qj in A[qi]:
                     logA[qi][qj] = np.log(A[qi][qj])
            self.A = logA

            # log(P)
            logP = {}
            for qi in P:
                logP[qi] = np.log(P[qi])
                
            self.P = logP
            self.B = B
            
    def a(self, q_1, q):
        return self.A[q_1][q]
    
    def b(self, q, x):
        likelihood_sum = 0.0
        target_prob = None

        for qi in self.B:
            (mu, sigma) = self.B[qi]
            
            prob = normpdf(x, mu, sigma)
            
            if qi == q:
                target_prob = prob + 0.00000001
                
            likelihood_sum += (prob + 0.00000001)
            
        return np.log(target_prob / likelihood_sum)
    
    def pi(self, q):
        return self.P[q]
    
    def state_num(self, q):
        return list(self.A.keys()).index(q)
    
    def num_state(self, num):
        return list(self.A.keys())[num]
    
    def viterbi(self, X):
        T = len(X)
        N = len(self.A)

        v_matrix = [[0 for _ in range(T)] for _ in range(N)]
        backtrack_matrix = [[-1 for _ in range(T)] for _ in range(N)]
    
        for t in range(T):
            x = X[t]
            
            if t == 0:
                for i in range(N):
                    q = self.num_state(i)
                    
                    transition_prob = self.pi(q)
                    emission_prob = self.b(q, x)
                    
                    prob = transition_prob + emission_prob
                    
                    v_matrix[i][t] = prob
                    
            else:
                for i in range(N):
                    
                    max_prob = float('-inf')
                    max_prop_state = -1
                    
                    q = self.num_state(i)
                    
                    emission_prob = self.b(q, x)
                    
                    for j in range(N):
                        q_1 = self.num_state(j)
                        
                        prev_prob = v_matrix[j][t - 1]
                        
                        transition_prob = self.a(q_1, q)
                        
                        prob = transition_prob + emission_prob + prev_prob
                        
                        if prob > max_prob:
                            max_prob = prob
                            max_prop_state = j
                            
                    v_matrix[i][t] = max_prob
                    backtrack_matrix[i][t] = max_prop_state
                    
        
        # Rekonstrukcija puta
        last_index = np.argmax(np.array(v_matrix)[:,t - 1])
        
        path = []
        t = T - 1
        
        while last_index != -1:
            last_state = self.num_state(last_index)
            path.append(last_state)
            
            last_index = backtrack_matrix[last_index][t]
            t -= 1
            
        return list(reversed(path))
    
    def forward(self, X, k = None):
        if k == None:
            T = len(X)
        else:
            T = k
        
        N = len(self.A)

        v_matrix = [[0 for _ in range(T)] for _ in range(N)]
        
        for t in range(T):
            x = X[t]
            
            if t == 0:
                for i in range(N):
                    q = self.num_state(i)
                    
                    transition_prob = self.pi(q)
                    emission_prob = self.b(q, x)

                    prob = transition_prob + emission_prob
                    
                    v_matrix[i][t] = prob
                    
            else:
                for i in range(N):

                    probs = []

                    q = self.num_state(i)

                    emission_prob = self.b(q, x)

                    for j in range(N):
                        q_1 = self.num_state(j)

                        prev_prob = v_matrix[j][t - 1]

                        transition_prob = self.a(q_1, q)

                        prob = transition_prob + emission_prob + prev_prob

                        probs.append(prob)

                    v_matrix[i][t] = logsumexp(probs)
                
        m = np.array(v_matrix)
        
        return logsumexp(m[:,T - 1]), m[:,T - 1], m
    
    def backward(self, X, k = None):
        T = len(X)
        N = len(self.A)
        
        if k == None:
            start = 0
        else:
            start = k

        v_matrix = [[0 for _ in range(T)] for _ in range(N)]
        
        for t in reversed(range(start, T)):
            x = X[t]
            
            if t == T - 1:
                for i in range(N):
                    q = self.num_state(i)
                    
                    transition_prob = np.log(1)
                    emission_prob = self.b(q, x)
                    
                    prob = transition_prob + emission_prob
                    
                    v_matrix[i][t] = prob
                    
            else:
                for i in range(N):

                    sum_prob = 0
                    probs = []

                    q = self.num_state(i)

                    emission_prob = self.b(q, x)

                    for j in range(N):
                        q_1 = self.num_state(j)

                        prev_prob = v_matrix[j][t + 1]

                        transition_prob = self.a(q, q_1)

                        prob = transition_prob + emission_prob + prev_prob
                        
                        if t == k:
                            prob += self.pi(q)

                        probs.append(prob)

                    v_matrix[i][t] = logsumexp(probs)
                
        m = np.array(v_matrix)
        
        return logsumexp(m[:,k]), m[:,k], m
    
    def baum_welch_single_sequence(self, X):
        _, _, all_alpha = self.forward(X)
        _, _, all_beta = self.backward(X)
        
        T = len(X)
        N = len(self.A)
        
        gamma = np.matrix([[0.0 for j in range(T)] for i in range(N)])
        
        for t in range(T):
            marg_probs = []
            
            for i in range(N):
                prob = logsumexp([all_alpha[i,t], all_beta[i,t]])
                marg_probs.append(prob)            
                gamma[i,t] = prob
                
            gamma[:, t] -= logsumexp(marg_probs)
            
        zeye = np.array([[[0.0 for t in range(T)] for j in range(N)] for i in range(N)])
        
        for t in range(T - 1):
            marg_probs = []
            for i in range(N):
                qi = self.num_state(i)
                
                for j in range(N):
                    qj = self.num_state(j)
                    
                    prob = all_alpha[i, t] + self.a(qi, qj) + all_beta[j, t + 1] + self.b(qj, X[t + 1])
                    zeye[i, j, t] = prob
                    marg_probs.append(prob)
                    
            zeye[:,:,t] -= logsumexp(marg_probs)
            
        new_P = {}
        
        for i in range(N):
            qi = self.num_state(i)
            
            new_P[qi] = gamma[i, 0]
            
        new_a = {}
        
        for i in range(N):
            qi = self.num_state(i)
            
            if qi not in new_a:
                new_a[qi] = {}
            
            for j in range(N):
                qj = self.num_state(j)
                
                new_a[qi][qj] = logsumexp(zeye[i,j, : T - 1]) - logsumexp(gamma[i, : T - 1])
    
                
        v = ['+','-','0']
        
        new_b = {}
    
        Xarr = np.array(X)
        
        for i in range(N):

            qi = self.num_state(i)
                
            mu = ((gamma[i,:].dot(X)) / gamma[i,:].sum()).item()

            S = (Xarr - mu) * (Xarr - mu)

            sigma = np.sqrt(((gamma[i,:].dot(S)) / gamma[i,:].sum()).item())

            new_b[qi] = (mu, sigma)
            
        return new_a, new_b, new_P
    
    def x_prob(self, X_arr):
        return self.forward(X_arr)[0]
    
    def baum_welch(self, X_arr):
        R = len(X_arr)
        N = len(self.A)
        
        eps = 0.0001
            
        old_prob = 0
        new_prob = 1
        
        v = ['+','-','0']

        it = 1
        
        while True:
            print(f'Iteration: {it}')
            old_prob = self.x_prob(X)
            print(f'start log(p) = {old_prob}')
            
            a = []
            b = []
            p = []
            
            for x in X_arr:
                ai, bi, pi = self.baum_welch_single_sequence(x)
                a.append(ai)
                b.append(bi)
                p.append(pi)
                
            new_P = {}
            new_A = {}
            new_B = {}
                
            for r in range(R):
                # P
                for i in range(N):
                    qi = self.num_state(i)
                    
                    if qi not in new_P:
                        new_P[qi] = 0
                        
                    new_P[qi] += (p[r][qi] / R)
                    
                    
                # A
                for i in range(N):
                    qi = self.num_state(i)
                    for j in range(N):
                        qj = self.num_state(j)

                        if qi not in new_A:
                            new_A[qi] = {}
                            
                        if qj not in new_A[qi]:
                            new_A[qi][qj] = 0

                        new_A[qi][qj] += (a[r][qi][qj] / R)
                        
                # B
                for i in range(N):
                    qi = self.num_state(i)

                    mu = 0
                    sigma = 0

                    mu += ((b[r][qi][0]) / (R*len(v)))
                    sigma += ((b[r][qi][1]) / (R*len(v)))

                    new_B[qi] = (mu, sigma)

            self.A = new_A
            self.P = new_P
            self.B = new_B
            
            new_prob = self.x_prob(X)
            print(f'log(p) = {new_prob}')

            it += 1
    
    def forward_backward(self, X, t):
        _, alpha, _ = self.forward(X, t)
        _, beta, _ = self.backward(X, t)
        prod = (alpha + beta)
        norm_prod = prod - logsumexp(prod)
        
        return norm_prod